import os
from CalTorque import CalTorque
import matplotlib
from matplotlib import pyplot as PLT
from datetime import datetime
from scipy.stats import linregress

data_files = []
dp = 'data/'
ls= os.listdir(dp)
for i in ls:
   if i.split('.')[-1] not in ['xlsx', 'xls']: continue
   if 'out' in i: continue
   elif 'in' not in i: continue
   data_files.append(dp+i)

parameters = {}
parameters['p0'] = []
parameters['p1'] = []
parameters['p2'] = []
parameters['p3'] = []
parameters['p4'] = []
parameters['p5'] = []
parameters['p6'] = []
parameters['p7'] = []
parameters['p8'] = []
filelist = []
date = []
intdate = []


for f in data_files:
    data = CalTorque(f)
    if data['source'] != 'weakTh-228': continue
    if data['station'] != 'S5': continue
    if max(data['pos_offsetfix']) < 350: continue   # because some stations in run_history are wrong
    if max(data['pos_offsetfix']) > 400: continue
    if 'new_parms' not in data.keys(): continue
    print f
    
    parameters['p0'].append(data['new_parms'][0])
    parameters['p1'].append(data['new_parms'][1])
    parameters['p2'].append(data['new_parms'][2])
    parameters['p3'].append(data['new_parms'][3])
    parameters['p4'].append(data['new_parms'][4])
    parameters['p5'].append(data['new_parms'][5])
    parameters['p6'].append(data['new_parms'][6])
    parameters['p7'].append(data['new_parms'][7])
    parameters['p8'].append(data['new_parms'][8])
    
    filelist.append(f.split('.')[0].split(' ')[2].split('_')[0])
    rundate = datetime.strptime(str(data['datetime'][0]).split(' ')[0], '%Y-%m-%d')
    date.append(rundate)


for d in date:
    intd = d - min(date)
    intdate.append(intd.days)


l0 = linregress(intdate, parameters['p0'])
l1 = linregress(intdate, parameters['p1'])
l2 = linregress(intdate, parameters['p2'])
l3 = linregress(intdate, parameters['p3'])
l4 = linregress(intdate, parameters['p4'])
l5 = linregress(intdate, parameters['p5'])
l6 = linregress(intdate, parameters['p6'])
l7 = linregress(intdate, parameters['p7'])
l8 = linregress(intdate, parameters['p8'])

linearfits = {}
linearfits['p0'] = []
linearfits['p1'] = []
linearfits['p2'] = []
linearfits['p3'] = []
linearfits['p4'] = []
linearfits['p5'] = []
linearfits['p6'] = []
linearfits['p7'] = []
linearfits['p8'] = []

def Polyfnc(x, m, b):
    return m*x + b

for thing in intdate:
    linearfits['p0'].append(Polyfnc(thing, l0[0], l0[1]))
    linearfits['p1'].append(Polyfnc(thing, l1[0], l1[1]))
    linearfits['p2'].append(Polyfnc(thing, l2[0], l2[1]))
    linearfits['p3'].append(Polyfnc(thing, l3[0], l3[1]))
    linearfits['p4'].append(Polyfnc(thing, l4[0], l4[1]))
    linearfits['p5'].append(Polyfnc(thing, l5[0], l5[1]))
    linearfits['p6'].append(Polyfnc(thing, l6[0], l6[1]))
    linearfits['p7'].append(Polyfnc(thing, l7[0], l7[1]))
    linearfits['p8'].append(Polyfnc(thing, l8[0], l8[1]))



## ____ Plots ____ ##

box = [0.14, 0.14, 0.76, 0.76]

markers = ['o', '^', 'D', 'o', 'D', 's', 'p', 'h', 'd'] 
colors = ['b', 'g', 'r', 'r', 'g', 'b', 'r', 'g', 'b']
ylabels = ['A6 Peak Amplitude (in.oz)', 'A6 Peak Location (cm)', 'Gaussian Width (cm)', 
'A7 Peak Amplitude (in.oz)', 'Gaussian Width (cm)', 'Polynomial Offset (cm)', 'Leading Coefficient', 
'First Power Coefficient', 'Constant Coefficient']
titles = ['A6 Peak Amplitude Versus Date', 'A6 Peak Location Versus Date', 
'Gaussian Width Versus Date for Peak at Bend A6', 'A7 Peak Amplitude (in.oz)', 
'Gaussian Width Versus Date for Peak at Bend A7', 'Polynomial Position Offset Versus Date', 
'Second Order Polynomial Leading Coefficient Versus Date', 
"Second Order Polynomial's First Power Coefficient Versus Date", 
"Second Order Polynomial's Constant Coefficient"]

for p in parameters:
    print list(p)
    fig = PLT.figure(figsize = (15, 8), dpi = 150)
    ax = fig.add_axes(box)
    ax.set_ylabel(ylabels[int(list(p)[-1])])
    ax.set_xlabel('Date')
    PLT.title(titles[int(list(p)[-1])])
    ax.grid()
    ax.scatter(date, parameters[p], color = colors[int(list(p)[-1])], marker = markers[int(list(p)[-1])])
    ax.plot(date, linearfits[p], '-', color = colors[int(list(p)[-1])])
    fig.savefig('parameters/' + p + '_vs_date.png')



fig = PLT.figure(figsize = (15, 8), dpi = 150)
axa = fig.add_axes(box)
axa.set_ylabel('Gaussian Width (cm)')
axa.set_xlabel('Date')
PLT.title('Gaussian Width Versus Date')
axa.grid()
axa.scatter(date, parameters['p2'], color = 'r', marker = 'D', label = 'Peak at Bend A6')
axa.plot(date, linearfits['p2'], '-', color = 'r')
axa.scatter(date, parameters['p4'], color = 'g', marker = 'D', label = 'Peak at Bend A7')
axa.plot(date, linearfits['p4'], '-', color = 'g')
h, l = axa.get_legend_handles_labels()
#for i, txt in enumerate(filelist):
    #axa.annotate(txt, (date[i], parameters['p4'][i]), position = (date[i], parameters['p4'][i]), size = 'x-small', alpha = 0.5, rotation = 45)
PLT.legend(h, l, 'upper left')
fig.savefig('parameters/p2_p4_vs_date.png')

figb = PLT.figure(figsize = (15, 8), dpi = 150)
axb = figb.add_axes(box)
axb.set_ylabel('Amplitude (in.oz)')
axb.set_xlabel('Date')
PLT.title('Amplitude of Peaks at A6 and A7 Versus Date')
axa.grid()
axb.scatter(date, parameters['p0'], color = 'b', marker = 'o', label = 'A6 Peak')
axb.plot(date, linearfits['p0'], '-', color = 'b')
axb.scatter(date, parameters['p3'], color = 'r', marker = 'o', label = 'A7 Peak')
axb.plot(date, linearfits['p3'], '-', color = 'r')
h, l = axb.get_legend_handles_labels()
PLT.legend(h, l, 'upper left')
figb.savefig('parameters/p0_p3_vs_date.png')


